#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <stdlib.h>

/*
 * y4m file form
 * +---------------------------+
 * | YUV4MPEG2 para1 para2 ... |
 * +---------------------------+
 * | FRAME                     |
 * | YUV DATA                  |
 * +---------------------------+
 * | FRAME                     |
 * | YUV DATA                  |
 * +---------------------------+
 * | --- ---                   |
 * | --- ---                   |
 * +---------------------------+
 * | FRAME                     |
 * | YUV DATA                  |
 * +---------------------------+
 * 明文开始：YUV4MPEG2
 * W：帧宽 比如 W720
 * H：帧高 比如 H480
 * F：帧率 比如 F30:1 30fps
 * I：交错模式 Ip 表示逐行 Progressive
 * A：像素宽高比 A1:1 square pixels 方形像素
 * C：色彩空间 C420 C422 C444等等
 * 参数信息后面保存原始数据YUV
 * 每帧都以FRAME开始，再在后面加一个0x0A
 * 之后就是原始的图像帧
 * 详细信息：https://blog.csdn.net/hiccupzhu/article/details/19498093
 * Y4M yuv格式为YU12 YYYYYYYY UU VV
 * NV12 格式为 YYYYYYYY UVUV
 * NV21 格式为 YYYYYYYY VUVU
 */

enum YUV_MODE {
	YUV_MODE_NV12,
	YUV_MODE_NV21
}yuv_mode;

int width = 0, height = 0;
int fps = 0;
static char optstr[] = "o:m:i:h";

static const struct option long_options[] = {
	{"input", required_argument, NULL, 'i'},
	{"output", optional_argument, NULL, 'o'},
	{"mode", optional_argument, NULL, 'm'},
	{"help", no_argument, NULL, 'h'}
};

void print_usage(void)
{
	printf("examples:\n");
	printf("\t[-i inputfile] [-m [mode]] [-o [outputfile]]\n");
	printf("\t-i | --input : inputfile y4m form\n");
	printf("\t-m | --mode : nv12 nv21, Default nv12\n");
	printf("\t-o | --output: output file, Default output.nv12\n");
}

#define Y4MHEAD 		"YUV4MPEG2 "
#define Y4MHEAD_LEN 	10
#define FRAME_FLAG  	"FRAME"
#define FRAME_FLAG_LEN 	5

int check_head(FILE * in_fp)
{
	int len = 0;
	char data_buf[12];
	int i,n = 1, p;

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: YUV4MPEG2 faile\n");
		return -1;
	}

	if(strncmp(data_buf, Y4MHEAD, Y4MHEAD_LEN)) {
		printf("match YUV4MPEG2 faile\n");
		return -1;
	}

	len += Y4MHEAD_LEN;

	fseek(in_fp, len, SEEK_SET);

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: width faile\n");
		return -1;
	}
	
	if (data_buf[0] == 'W') {
		for(i=1; i<12; i++) {
			if (data_buf[i] == ' ')
				break;
		}
		len += i+1;
		i -= 1;
		n = 1;
		for(i; i>0; i--) {
			width += (data_buf[i] - '0')*n;
			n *= 10;
		}
	} else {
		printf("match width faile\n");
		return -1;
	}

	printf("width = %d\n", width);
	fseek(in_fp, len, SEEK_SET);

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: height faile\n");
		return -1;
	}

	if (data_buf[0] == 'H') {
		for(i=1; i<12; i++) {
			if (data_buf[i] == ' ')
				break;
		}
		len += i+1;
		i -= 1;
		n = 1;
		for(i; i>0; i--) {
			height += (data_buf[i] - '0')*n;
			n *= 10;
		}
	} else {
		printf("match height faile\n");
		return -1;
	}

	printf("height = %d\n", height);
	fseek(in_fp, len, SEEK_SET);

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: fps faile\n");
		return -1;
	}

	if (data_buf[0] == 'F') {
		for(i=1; i<12; i++) {
			if (data_buf[i] == ':')
				p = i;
			if (data_buf[i] == ' ')
				break;
		}
		len += i+1;
		p -= 1;
		n = 1;
		for(p; p>0; p--) {
			fps += (data_buf[p] - '0')*n;
			n *= 10;
		}
	} else {
		printf("match fps faile\n");
		return -1;
	}
	printf("fps = %d\n", fps);

	fseek(in_fp, len, SEEK_SET);

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: I & A faile\n");
		return -1;
	}

	//ignore I & A
	for(i = 0, n = 0; i < 12, n !=2; i++) {
		if((data_buf[i] == ' ') || (data_buf[i] == '\n'))
			n++;
	}

	len += i;
	fseek(in_fp, len, SEEK_SET);

	memset(data_buf, 0, 12);
	if(fread(data_buf, 1, 12, in_fp) != 12) {
		printf("get head: I & A faile\n");
		return -1;
	}

	printf("data_buf[0] = %c\n", data_buf[0]);
	return len;
}

int main(int argc, char **argv)
{
	char *in = "NULL";
	char *out = "output.nv12";
	enum YUV_MODE mode = YUV_MODE_NV12;
	int frame_offset, frame_cnt, frame_size;
	char *frame_buf;
	char *uv_frame_buf;
	FILE* in_fp;
	FILE* out_fp;

	int c;
	while ((c = getopt_long(argc, argv, optstr, long_options, NULL)) != -1) {
		const char *tmp_optarg = optarg;
		switch (c) {
			case 'i':
				if (tmp_optarg) {
					in = (char *)tmp_optarg;
				} else {
					printf("input file error!\n");
					return -1;
				}
			break;
			case 'm':
				if (!strcmp(optarg, "nv12")) {
					mode = YUV_MODE_NV12;
				} else if (!strcmp(optarg, "nv21")) {
					mode = YUV_MODE_NV21;
				} else {
					printf("ERROR: Unsupport yuv mode\n");
					return -1;
				}
			break;
			case 'o':
				if (tmp_optarg) {
					out = (char *)tmp_optarg;
				} else {
					printf("output file error!\n");
					return -1;
				}
			break;
			case 'h':
			default:
				print_usage();
				return -1;
		}
	}

	printf("input file : %s\n", in);
	printf("output file : %s\n", out);
	printf("yuv mode : %d\n", mode);
	
	in_fp = fopen(in, "r");
	if (!in_fp) {
		printf("open %s failed\n", in);
		return -1;
	}
	out_fp = fopen(out, "w");
	if (!out_fp) {
		printf("open %s failed\n", out);
		return -1;
	}

	frame_offset = check_head(in_fp);
	if (frame_offset < 0) {
		printf("get offset of framedata fail\n");
		return -1;
	}

	fseek(in_fp, frame_offset, SEEK_SET);
	frame_size = width * height * 15 / 10 + 6;
	frame_buf = (char*)malloc(frame_size);
	if (!frame_buf) {
		printf("alloc framebuff fail\n");
		return -1;
	}
	frame_cnt = 0;
	int y_size = width*height;
	int u_offset = width*height;
	int u_v_size = width*height/4;
	int v_offset = width*height + u_v_size;
	int i;
	printf("u_offset = %d, v_offset = %d, u_v_size = %d\n", u_offset, v_offset, u_v_size);
	uv_frame_buf = (char*)malloc(u_v_size * 2);
	if (!uv_frame_buf) {
		printf("alloc uv framebuff fail\n");
		fclose(in_fp);
		free(frame_buf);
		return -1;
	}

	while(fread(frame_buf, 1, frame_size, in_fp) == frame_size) {
		if(strncmp(frame_buf, FRAME_FLAG, FRAME_FLAG_LEN)) {
			printf("invalid frame data\n");
			continue;
		}
		frame_cnt ++;

		if (mode == YUV_MODE_NV12) {
			//YU12 -> NV12
			for(i=0; i<u_v_size; i++) {
				uv_frame_buf[i*2] = frame_buf[FRAME_FLAG_LEN + 1 + u_offset + i];
				uv_frame_buf[i*2 + 1] = frame_buf[FRAME_FLAG_LEN + 1 + v_offset + i];
			}
		} else if(mode == YUV_MODE_NV21) {
			//YU12 -> NV21
			for(i=0; i<u_v_size; i++) {
				uv_frame_buf[i * 2] = frame_buf[FRAME_FLAG_LEN + 1 + v_offset + i];
				uv_frame_buf[i * 2 + 1] = frame_buf[FRAME_FLAG_LEN + 1 + u_offset + i];
			}
		}
		memcpy(&frame_buf[FRAME_FLAG_LEN + 1 + y_size], uv_frame_buf, u_v_size * 2);

		fwrite(&frame_buf[FRAME_FLAG_LEN + 1], 1, frame_size - (FRAME_FLAG_LEN + 1 ), out_fp);
	}

	printf("frame count = %d\n", frame_cnt);
	
	fclose(in_fp);
	fclose(out_fp);
	free(frame_buf);
	free(uv_frame_buf);
	return 0;
}
